package ms.unicorn.dsl.eistar.reasoning.maxbelief;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import ms.unicorn.dsl.eistar.util.Probability;
import ms.unicorn.utils.Console;
import eistar.Entity;
import eistar.Goal;
import eistar.Model;
import eistar.Relation;

public class MaxBeliefTable {
	
	public static class Solution {
		private Goal topGoal;
		private Set<Goal> configuration;
		
		public Solution() {
			
		}
		
		public Solution(Goal topGoal, Set<Goal> configuration) {
			this.topGoal = topGoal;
			this.configuration = configuration;
		}
		
		public Goal getTopGoal() {
			return topGoal;
		}
		
		public void setTopGoal(Goal topGoal) {
			this.topGoal = topGoal;
		}

		public Set<Goal> getConfiguration() {
			return configuration;
		}

		public void setConfiguration(Set<Goal> configuration) {
			this.configuration = configuration;
		}
		
		@Override
		public boolean equals(Object obj) {
			if (obj instanceof Solution)
				return equals((Solution)obj);
			return super.equals(obj);
		}
		
		public boolean equals(Solution sol) {
			boolean equal = false;
			if (sol != null) {
				equal = topGoal != null && sol.getTopGoal() != null 
						&& configuration != null  && sol.getConfiguration() != null
						&& topGoal.equals(sol.getTopGoal()) 
						&& configuration.equals(sol.getConfiguration());
			}
			return equal;
		}
		
		public boolean subsumes(Solution sol) {
			boolean subsumed = topGoal != null && sol.getTopGoal() != null 
					&& configuration != null  && sol.getConfiguration() != null
					&& configuration.containsAll(sol.getConfiguration());
			return subsumed;
		}
	
		@Override
		public String toString() {
			return String.format("<%s, %s>", String.valueOf(topGoal), String.valueOf(configuration));
		}
	}
	
	public static class MaxBeliefEntry  {
		private Solution solution;
		private int maxBelief, residualRisk;
		
		public MaxBeliefEntry(Goal goal, Set<Goal> configuration) {
			this(new Solution(goal, configuration), Probability.ONE, 0);
		}
		
		public MaxBeliefEntry(Solution solution, int maxBelief, int residualRisk) {
			this.solution = solution;
			this.maxBelief = maxBelief;
			this.residualRisk = residualRisk;
		}
		
		@Override
		public boolean equals(Object obj) {
			if (obj instanceof MaxBeliefEntry) 
				return equals((MaxBeliefEntry)obj);
			return super.equals(obj);
		}
		
		public boolean equals(MaxBeliefEntry entry) {
			boolean equal = false;
			if (entry != null)
				equal = solution.equals(entry.getSolution());
			return equal;
		}

		public boolean subsumes(MaxBeliefEntry entry) {
			boolean subsumed = false;
			if (entry != null)
				subsumed = solution.subsumes(entry.solution);
			return subsumed;
		}
		
		public void setResidualRisk(int residualRisk) {
			this.residualRisk = residualRisk;
		}

		public int getResidualRisk() {
			return residualRisk;
		}

		public void setMaxBelief(int maxBelief) {
			this.maxBelief = maxBelief;
		}

		public int getMaxBelief() {
			return maxBelief;
		}
	
		@Override
		public String toString() {
			return String.format("%s, %d, %d", solution, maxBelief, residualRisk);
		}
	
		public MaxBeliefEntry clone() {
			MaxBeliefEntry result = new MaxBeliefEntry(solution.getTopGoal(), 
					new TreeSet<Goal>(solution.getConfiguration()));
			result.maxBelief = maxBelief;
			result.residualRisk = residualRisk;
			return result;
		}
		
		public void mergeProbability(MaxBeliefEntry t) {
			setMaxBelief(Math.max(maxBelief, t.maxBelief));
			setResidualRisk(Math.max(0, residualRisk + t.residualRisk - Probability.ONE));
		}
		
		public void joinProbability(MaxBeliefEntry t) {
			setMaxBelief(Probability.mult(maxBelief, t.maxBelief));
			setResidualRisk(Probability.ONE - Probability.mult(Probability.ONE - residualRisk, Probability.ONE - t.residualRisk));
		}
		
		public Solution getSolution() {
			return solution;
		}

		public void setSolution(Solution solution) {
			this.solution = solution;
		}
	}

	private List<MaxBeliefEntry> entries;
	private Set<Goal> topGoals;

	
	public MaxBeliefTable() {
		entries = new ArrayList<MaxBeliefTable.MaxBeliefEntry>();
	}

	public List<MaxBeliefEntry> getEntries() {
		return entries;
	}
	
	public Set<Goal> getTopGoals() {
		if (topGoals == null) {
			topGoals = new TreeSet<Goal>();
			for(MaxBeliefEntry e: entries)
				topGoals.add(e.solution.getTopGoal());
		}
		return topGoals;
	}
	
	public List<MaxBeliefEntry> getEntriesOfGoal(Goal topGoal) {
		ArrayList<MaxBeliefEntry> list = new ArrayList<MaxBeliefEntry>();
		for(MaxBeliefEntry e: entries)
			if (e.getSolution().getTopGoal().equals(topGoal))
				list.add(e);
		return list;
	}
	
	public void generateFromModel(Model m) {
		generateFromModel(m, null);
	}
	
	public void generateFromModel(Model m, List<MaxBeliefTable> subtables) {
		clear();
		//ArrayList<MaxBeliefEntry> pendingList = new ArrayList<MaxBeliefTable.MaxBeliefEntry>();
		Map<Goal, List<MaxBeliefEntry>> entries = new HashMap<Goal, List<MaxBeliefEntry>>();
			
		Set<Entity> cutPoints =  null;
		topGoals = new TreeSet<Goal>();
		
		// 1st round: initialize the table
		// a) append cut-points
		if (subtables != null) {
			cutPoints = new TreeSet<Entity>();
			for (MaxBeliefTable t: subtables) 
				for(MaxBeliefEntry e: t.entries) {
					Goal pGoal = e.solution.topGoal;
					List<MaxBeliefEntry> list = entries.get(pGoal);
					if (list == null)
						entries.put(pGoal, list = new ArrayList<MaxBeliefTable.MaxBeliefEntry>());
					list.add(e.clone());
					topGoals.add(pGoal);
					cutPoints.add(pGoal);
					
				}
		}
		// b) append non-cut point decomposition and delegation
		for(Relation rel: m.getModelElement(Relation.class))
			switch(rel.getRelType()) {
			case DECOMPOSE:
			case DELEGATE:
				addRelationToTable(rel, cutPoints, entries);
			}
		
				
	
		// 2nd round: go through the table and expand the relation, loop until their is no expandable rule
		// In this round we use MaxBelief of each entry as a flag which could be one of the following
		//		0 		: the goal of this entry is a top goal
		//		MAX_INT : the goal of this entry is not a top goal
		boolean conti = false;
		do { 
			conti = false;
			for(List<MaxBeliefEntry> list: entries.values()) 
				for(int i = 0; i < list.size(); i++) {
					MaxBeliefEntry entry = list.get(i);
					if (topGoals.contains(entry.solution.topGoal)) { // only process top goal
						for(Object go: entry.getSolution().configuration.toArray()) { // walk through all goal in this configuration
							Goal g = (Goal)go;
							// check for decomposition of goal g.
							if (entries.containsKey(g)) { 
								conti = true; // we need another loop
								List<MaxBeliefEntry> subrules = entries.get(g);
								expandEntry(list, entry, g, subrules);
							}
						}
					}
				}
		} while (conti);
		
		
		// 3rd round: append top-goal entries to the table entries
		for (Goal goal : entries.keySet()) 
			if (topGoals.contains(goal)) {
				List<MaxBeliefEntry> list = entries.get(goal);
				for (MaxBeliefEntry entry: list)
					this.entries.add(entry);
			}
	}
	
	private void expandEntry(List<MaxBeliefEntry> list, MaxBeliefEntry entry, Goal g, List<MaxBeliefEntry> subEntries) {
		boolean first = true;
		MaxBeliefEntry orgEntry = subEntries.size() > 1 ? entry.clone() : entry;
		for (MaxBeliefEntry subentry : subEntries) {
			//exclude the goal of subentry from TopGoal
			topGoals.remove(subentry.solution.topGoal);
			MaxBeliefEntry newEntry = first ? entry : orgEntry.clone();
			if (!first) 
				list.add(newEntry);
			
			newEntry.solution.configuration.remove(g);
			//for(Goal sg: subentry.solution.configuration)
			//	newEntry.solution.configuration.add(sg);
			newEntry.solution.configuration.addAll(subentry.solution.configuration);
			
			// update the max belief && residual risk
			newEntry.joinProbability(subentry);
			
			first = false;
		}
	}

	private void addRelationToTable(Relation rel, Set<Entity> cutPoints, Map<Goal, List<MaxBeliefEntry>> entries) {
		Goal pGoal = (Goal) rel.getSource();
		if (cutPoints == null || !cutPoints.contains(pGoal)) {
			topGoals.add(pGoal);
			Set<Goal> cGoals = new TreeSet<Goal>();
			for (Entity target: rel.getTarget())
				cGoals.add((Goal) target);
			
			List<MaxBeliefEntry> list = entries.get(pGoal);
			if (list == null)
				entries.put(pGoal, list = new ArrayList<MaxBeliefTable.MaxBeliefEntry>());
			list.add(new MaxBeliefEntry(pGoal, cGoals));
		}
		else if (cutPoints != null)
			cutPoints.addAll(rel.getTarget());
	}

	public void printTableContent() {
		for(MaxBeliefEntry entry: entries)
			Console.println(entry.toString());
	}

	public void clear() {
		entries.clear();	
	}

	public void resetMaxBelief(int value) {
		for(MaxBeliefEntry entry: entries)
			entry.setMaxBelief(value);
	}
	
	public void resetResidualRisk(int value) {
		for(MaxBeliefEntry entry: entries)
			entry.setResidualRisk(value);
	}
	
	public static MaxBeliefTable sequentialMerge(MaxBeliefTable ltable, MaxBeliefTable rtable) {
		MaxBeliefTable table = new MaxBeliefTable();
		for(MaxBeliefEntry e: ltable.entries)
			table.mergeOneEntry(e, rtable, true);
		for(MaxBeliefEntry e: rtable.entries)
			table.mergeOneEntry(e, ltable, false);
		return table;
	}
	
	protected void mergeOneEntry(MaxBeliefEntry entry, MaxBeliefTable sourceTable, boolean addEntryOnEqual) {
		MaxBeliefEntry newEntry = entry.clone();
		for(MaxBeliefEntry sEntry: sourceTable.entries) {
			if (newEntry.equals(sEntry))
				if (!addEntryOnEqual) return;
				else 
					newEntry.mergeProbability(sEntry);
			else if (newEntry.subsumes(sEntry))
				newEntry.mergeProbability(sEntry);
				
		}
		entries.add(newEntry);
	}
	
	public MaxBeliefTable merge(MaxBeliefTable btable) {
		return sequentialMerge(this, btable);
	}
}
